#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""  
@Project : PythonProject
@File : ssh_val.py
@Author : donc
@Time : 2025/5/12 19:10
@脚本说明 : SSH弱口令检测模块
"""
import threading
import paramiko
from queue import Queue
import socket
import time


def test_ssh_connection(host, port):
    """测试SSH服务是否可达"""
    try:
        # 创建TCP socket
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(3)

        # 尝试连接
        result = sock.connect_ex((host, int(port)))
        sock.close()

        return result == 0
    except Exception as e:
        return False


def val_ssh(host, port, username_file, password_file, progress_callback=None):
    """
    SSH弱口令检测主函数
    返回格式: {"host": [], "port": [], "username": [], "password": []}
    """
    # 读取用户名字典
    try:
        usernames = username_file.getvalue().decode().splitlines()
    except Exception as e:
        return {"error": f"用户名字典读取失败: {str(e)}"}

    # 读取密码字典
    try:
        passwords = password_file.getvalue().decode().splitlines()
    except Exception as e:
        return {"error": f"密码字典读取失败: {str(e)}"}

    # 计算总凭证数
    total_credentials = len(usernames) * len(passwords)

    # 初始化进度
    if progress_callback:
        progress_callback(0, total_credentials)

    # 创建线程安全队列
    result_queue = Queue()
    threads = []
    processed = 0  # 跟踪已处理的凭证数

    # 创建并启动线程
    for username in usernames:
        for password in passwords:
            t = threading.Thread(
                target=try_ssh_connection,
                args=(host, int(port), username.strip(), password.strip(), result_queue)
            )
            t.start()
            threads.append(t)

    # 等待所有线程完成并更新进度
    while threads:
        # 收集已完成的结果
        for t in threads[:]:
            if not t.is_alive():
                threads.remove(t)
                processed += 1

                # 更新进度
                if progress_callback:
                    progress_callback(processed, total_credentials)

        # 收集队列中的结果
        while not result_queue.empty():
            result_queue.get()  # 实际结果已经在try_ssh_connection中处理
            processed += 1
            if progress_callback:
                progress_callback(processed, total_credentials)

        # 稍微等待一下
        time.sleep(0.1)

    # 确保最终进度是100%
    if progress_callback:
        progress_callback(total_credentials, total_credentials)

    # 收集所有结果
    results = {
        "host": [],
        "port": [],
        "username": [],
        "password": []
    }

    # 注意：在实际应用中，结果已经在try_ssh_connection中处理
    # 这里返回空结果，因为实际结果已经在回调中处理
    return results

def try_ssh_connection(host, port, username, password, result_queue):
    """尝试连接SSH并验证密码"""
    try:
        # 创建SSH客户端
        ssh_client = paramiko.SSHClient()

        # 自动添加主机密钥（生产环境应谨慎）
        ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        # 设置连接超时
        ssh_client.connect(
            hostname=host,
            port=port,
            username=username,
            password=password,
            timeout=5,
            banner_timeout=30
        )

        # 简单验证连接
        stdin, stdout, stderr = ssh_client.exec_command('echo "SSH Connection Successful"', timeout=5)
        output = stdout.read().decode().strip()

        # 如果连接成功，保存结果
        if "SSH Connection Successful" in output:
            # 将成功结果放入队列
            result_queue.put({
                "host": host,
                "port": port,
                "username": username,
                "password": password
            })

        # 关闭连接
        ssh_client.close()

    except paramiko.AuthenticationException:
        # 认证失败，忽略
        pass
    except paramiko.SSHException as e:
        # SSH协议错误
        pass
    except socket.timeout:
        # 连接超时
        pass
    except Exception as e:
        # 其他错误
        pass